# This file is part of Buildbot. Buildbot is free software: you can
# redistribute it and/or modify it under the terms of the GNU General Public
# License as published by the Free Software Foundation, version 2.
#
# This program is distributed in the hope that it will be useful, but WITHOUT
# ANY WARRANTY; without even the implied warranty of MERCHANTABILITY or FITNESS
# FOR A PARTICULAR PURPOSE. See the GNU General Public License for more
# details.
#
# You should have received a copy of the GNU General Public License along with
# this program; if not, write to the Free Software Foundation, Inc., 51
# Franklin Street, Fifth Floor, Boston, MA 02110-1301 USA.
#
# Copyright Buildbot Team Members

from __future__ import annotations

import os
import stat
from pathlib import PurePath
from typing import TYPE_CHECKING
from typing import Any
from unittest import mock

from parameterized import parameterized
from twisted.internet import defer
from twisted.trial import unittest

from buildbot.process import remotecommand
from buildbot.test.fake import fakemaster
from buildbot.test.reactor import TestReactorMixin
from buildbot.test.util import protocols as util_protocols
from buildbot.worker.protocols import base
from buildbot.worker.protocols import msgpack

if TYPE_CHECKING:
    from twisted.internet.defer import Deferred

    from buildbot.util.twisted import InlineCallbacksType


class TestListener(TestReactorMixin, unittest.TestCase):
    @defer.inlineCallbacks
    def setUp(self) -> InlineCallbacksType[None]:  # type: ignore[override]
        self.setup_test_reactor()
        self.master = yield fakemaster.make_master(self)

    def test_constructor(self) -> None:
        listener = msgpack.Listener(self.master)
        self.assertEqual(listener.master, self.master)
        self.assertEqual(listener._registrations, {})

    @defer.inlineCallbacks
    def test_update_registration_simple(self) -> InlineCallbacksType[None]:
        listener = msgpack.Listener(self.master)
        reg = yield listener.updateRegistration('example', 'pass', 'tcp:1234')
        self.assertEqual(self.master.msgmanager._registrations, [('tcp:1234', 'example', 'pass')])
        self.assertEqual(listener._registrations['example'], ('pass', 'tcp:1234', reg))

    @defer.inlineCallbacks
    def test_update_registration_pass_changed(self) -> InlineCallbacksType[None]:
        listener = msgpack.Listener(self.master)
        listener.updateRegistration('example', 'pass', 'tcp:1234')
        reg1 = yield listener.updateRegistration('example', 'pass1', 'tcp:1234')
        self.assertEqual(listener._registrations['example'], ('pass1', 'tcp:1234', reg1))
        self.assertEqual(self.master.msgmanager._unregistrations, [('tcp:1234', 'example')])

    @defer.inlineCallbacks
    def test_update_registration_port_changed(self) -> InlineCallbacksType[None]:
        listener = msgpack.Listener(self.master)
        listener.updateRegistration('example', 'pass', 'tcp:1234')
        reg1 = yield listener.updateRegistration('example', 'pass', 'tcp:4321')
        self.assertEqual(listener._registrations['example'], ('pass', 'tcp:4321', reg1))
        self.assertEqual(self.master.msgmanager._unregistrations, [('tcp:1234', 'example')])

    @defer.inlineCallbacks
    def test_create_connection(self) -> InlineCallbacksType[None]:
        listener = msgpack.Listener(self.master)
        listener.before_connection_setup = mock.Mock()  # type: ignore[method-assign]
        worker = mock.Mock()
        worker.workername = 'test'
        protocol = mock.Mock()

        listener.updateRegistration('example', 'pass', 'tcp:1234')
        self.master.workers.register(worker)
        conn = yield listener._create_connection(protocol, worker.workername)

        listener.before_connection_setup.assert_called_once_with(protocol, worker.workername)
        self.assertIsInstance(conn, msgpack.Connection)


class TestConnectionApi(
    util_protocols.ConnectionInterfaceTest, TestReactorMixin, unittest.TestCase
):
    @defer.inlineCallbacks
    def setUp(self) -> InlineCallbacksType[None]:  # type: ignore[override]
        self.setup_test_reactor()
        self.master = yield fakemaster.make_master(self)
        self.conn = msgpack.Connection(self.master, mock.Mock(), mock.Mock())


class TestConnection(TestReactorMixin, unittest.TestCase):
    @defer.inlineCallbacks
    def setUp(self) -> InlineCallbacksType[None]:  # type: ignore[override]
        self.setup_test_reactor()
        self.master = yield fakemaster.make_master(self)
        self.protocol = mock.Mock()
        self.worker = mock.Mock()
        self.worker.workername = 'test_worker'
        self.conn = msgpack.Connection(self.master, self.worker, self.protocol)

    def test_constructor(self) -> None:
        self.assertEqual(self.conn.protocol, self.protocol)
        self.assertEqual(self.conn.master, self.master)
        self.assertEqual(self.conn.worker, self.worker)

    @defer.inlineCallbacks
    def test_attached(self) -> InlineCallbacksType[None]:
        self.conn.attached(self.protocol)
        self.worker.attached.assert_called_with(self.conn)

        self.reactor.pump([10] * 361)
        self.protocol.get_message_result.assert_called_once_with({'op': 'keepalive'})

        self.conn.detached(self.protocol)
        yield self.conn.waitShutdown()

    @defer.inlineCallbacks
    def test_detached(self) -> InlineCallbacksType[None]:
        self.conn.attached(self.protocol)
        self.conn.detached(self.protocol)

        self.assertEqual(self.conn.keepalive_timer, None)
        self.assertEqual(self.conn.protocol, None)
        yield self.conn.waitShutdown()

    def test_lose_connection(self) -> None:
        self.conn.loseConnection()

        self.assertEqual(self.conn.keepalive_timer, None)
        self.protocol.transport.abortConnection.assert_called()

    def test_do_keepalive(self) -> None:
        self.conn._do_keepalive()
        self.protocol.get_message_result.assert_called_once_with({'op': 'keepalive'})

    @defer.inlineCallbacks
    def test_start_stop_keepalive_timer(self) -> InlineCallbacksType[None]:
        self.conn.startKeepaliveTimer()

        self.protocol.get_message_result.assert_not_called()

        self.reactor.pump([10] * 361)
        expected_call = [
            mock.call({'op': 'keepalive'}),
        ]
        self.assertEqual(self.protocol.get_message_result.call_args_list, expected_call)

        self.reactor.pump([10] * 361)
        expected_calls = [
            mock.call({'op': 'keepalive'}),
            mock.call({'op': 'keepalive'}),
        ]
        self.assertEqual(self.protocol.get_message_result.call_args_list, expected_calls)

        self.conn.stopKeepaliveTimer()

        self.reactor.pump([10] * 361)
        expected_calls = [
            mock.call({'op': 'keepalive'}),
            mock.call({'op': 'keepalive'}),
        ]
        self.assertEqual(self.protocol.get_message_result.call_args_list, expected_calls)

        yield self.conn.waitShutdown()

    @defer.inlineCallbacks
    def test_remote_keepalive(self) -> InlineCallbacksType[None]:
        yield self.conn.remoteKeepalive()
        self.protocol.get_message_result.assert_called_once_with({'op': 'keepalive'})

    @defer.inlineCallbacks
    def test_remote_print(self) -> InlineCallbacksType[None]:
        yield self.conn.remotePrint(message='test')
        self.protocol.get_message_result.assert_called_once_with({'op': 'print', 'message': 'test'})

    @defer.inlineCallbacks
    def test_remote_get_worker_info(self) -> InlineCallbacksType[None]:
        self.protocol.get_message_result.return_value = defer.succeed({'system': 'posix'})
        result = yield self.conn.remoteGetWorkerInfo()

        self.protocol.get_message_result.assert_called_once_with({'op': 'get_worker_info'})
        self.assertEqual(result, {'system': 'posix'})

    def set_up_set_builder_list(
        self,
        builders: list[tuple[str, str]],
        delete_leftover_dirs: bool = True,
    ) -> Deferred[list[str]]:
        self.protocol.command_id_to_command_map = {}

        def get_message_result(*args: Any) -> Deferred[Any]:
            d: Deferred[Any] = defer.Deferred()
            self.d_get_message_result = d
            return d

        self.protocol.get_message_result.side_effect = get_message_result
        self.conn.info = {'basedir': 'testdir'}
        self.conn.info['delete_leftover_dirs'] = delete_leftover_dirs
        self.conn.path_module = os.path
        self.conn.path_cls = PurePath
        d = self.conn.remoteSetBuilderList(builders)
        return d

    async def check_message_send_response(
        self,
        command_name: str,
        args: dict[str, list[str] | str],
        update_msg: list[tuple[str, Any]],
    ) -> None:
        command_id = remotecommand.RemoteCommand.get_last_generated_command_id()
        self.protocol.get_message_result.assert_called_once_with({
            'op': 'start_command',
            'command_id': command_id,
            'command_name': command_name,
            'args': args,
        })
        self.protocol.get_message_result.reset_mock()
        self.d_get_message_result.callback(None)

        remote_command = self.protocol.command_id_to_command_map[command_id]
        await remote_command.remote_update_msgpack(update_msg)
        await remote_command.remote_complete(None)

    def check_message_set_worker_settings(self) -> None:
        newline_re = r'(\r\n|\r(?=.)|\033\[u|\033\[[0-9]+;[0-9]+[Hf]|\033\[2J|\x08+)'
        self.protocol.get_message_result.assert_called_once_with({
            'op': 'set_worker_settings',
            'args': {
                'newline_re': newline_re,
                'max_line_length': 4096,
                'buffer_timeout': 5,
                'buffer_size': 64 * 1024,
            },
        })
        self.protocol.get_message_result.reset_mock()
        self.d_get_message_result.callback(None)

    @defer.inlineCallbacks
    def test_remote_set_builder_list_no_rmdir(self) -> InlineCallbacksType[None]:
        d = self.set_up_set_builder_list([('builder1', 'test_dir1'), ('builder2', 'test_dir2')])
        self.check_message_set_worker_settings()

        yield self.check_message_send_response(
            'listdir', {'path': 'testdir'}, [('files', ['dir1', 'dir2', 'dir3']), ('rc', 0)]
        )

        path = os.path.join('testdir', 'dir1')
        yield self.check_message_send_response('stat', {'path': path}, [('stat', (1,)), ('rc', 0)])

        path = os.path.join('testdir', 'dir2')
        yield self.check_message_send_response('stat', {'path': path}, [('stat', (1,)), ('rc', 0)])

        path = os.path.join('testdir', 'dir3')
        yield self.check_message_send_response('stat', {'path': path}, [('stat', (1,)), ('rc', 0)])

        paths = [
            os.path.join('testdir', 'info'),
            os.path.join('testdir', 'test_dir1'),
            os.path.join('testdir', 'test_dir2'),
        ]
        yield self.check_message_send_response('mkdir', {'paths': paths}, [('rc', 0)])

        r = yield d
        self.assertEqual(r, ['builder1', 'builder2'])
        self.protocol.get_message_result.assert_not_called()

    @defer.inlineCallbacks
    def test_remote_set_builder_list_do_rmdir(self) -> InlineCallbacksType[None]:
        d = self.set_up_set_builder_list([('builder1', 'test_dir1'), ('builder2', 'test_dir2')])
        self.check_message_set_worker_settings()

        yield self.check_message_send_response(
            'listdir', {'path': 'testdir'}, [('files', ['dir1', 'dir2', 'dir3']), ('rc', 0)]
        )

        path = os.path.join('testdir', 'dir1')
        yield self.check_message_send_response(
            'stat', {'path': path}, [('stat', (stat.S_IFDIR,)), ('rc', 0)]
        )

        path = os.path.join('testdir', 'dir2')
        yield self.check_message_send_response(
            'stat', {'path': path}, [('stat', (stat.S_IFDIR,)), ('rc', 0)]
        )

        path = os.path.join('testdir', 'dir3')
        yield self.check_message_send_response(
            'stat', {'path': path}, [('stat', (stat.S_IFDIR,)), ('rc', 0)]
        )

        paths = [
            os.path.join('testdir', 'dir1'),
            os.path.join('testdir', 'dir2'),
            os.path.join('testdir', 'dir3'),
        ]
        yield self.check_message_send_response('rmdir', {'paths': paths}, [('rc', 0)])

        paths = [
            os.path.join('testdir', 'info'),
            os.path.join('testdir', 'test_dir1'),
            os.path.join('testdir', 'test_dir2'),
        ]
        yield self.check_message_send_response('mkdir', {'paths': paths}, [('rc', 0)])

        r = yield d
        self.assertEqual(r, ['builder1', 'builder2'])
        self.protocol.get_message_result.assert_not_called()

    @defer.inlineCallbacks
    def test_remote_set_builder_list_no_rmdir_leave_leftover_dirs(
        self,
    ) -> InlineCallbacksType[None]:
        d = self.set_up_set_builder_list(
            [('builder1', 'test_dir1'), ('builder2', 'test_dir2')], delete_leftover_dirs=False
        )
        self.check_message_set_worker_settings()

        yield self.check_message_send_response(
            'listdir', {'path': 'testdir'}, [('files', ['dir1', 'dir2', 'dir3']), ('rc', 0)]
        )

        paths = [
            os.path.join('testdir', 'info'),
            os.path.join('testdir', 'test_dir1'),
            os.path.join('testdir', 'test_dir2'),
        ]
        yield self.check_message_send_response('mkdir', {'paths': paths}, [('rc', 0)])

        r = yield d
        self.assertEqual(r, ['builder1', 'builder2'])
        self.protocol.get_message_result.assert_not_called()

    @defer.inlineCallbacks
    def test_remote_set_builder_list_no_mkdir_from_files(self) -> InlineCallbacksType[None]:
        d = self.set_up_set_builder_list([('builder1', 'test_dir1'), ('builder2', 'test_dir2')])
        self.check_message_set_worker_settings()

        yield self.check_message_send_response(
            'listdir', {'path': 'testdir'}, [('files', ['dir1', 'test_dir2']), ('rc', 0)]
        )

        path = os.path.join('testdir', 'dir1')
        yield self.check_message_send_response('stat', {'path': path}, [('stat', (1,)), ('rc', 0)])

        paths = [os.path.join('testdir', 'info'), os.path.join('testdir', 'test_dir1')]
        yield self.check_message_send_response('mkdir', {'paths': paths}, [('rc', 0)])

        r = yield d
        self.assertEqual(r, ['builder1', 'builder2'])
        self.protocol.get_message_result.assert_not_called()

    @defer.inlineCallbacks
    def test_remote_set_builder_list_no_mkdir(self) -> InlineCallbacksType[None]:
        d = self.set_up_set_builder_list([('builder1', 'test_dir1'), ('builder2', 'test_dir2')])
        self.check_message_set_worker_settings()

        yield self.check_message_send_response(
            'listdir',
            {'path': 'testdir'},
            [('files', ['test_dir1', 'test_dir2', 'info']), ('rc', 0)],
        )

        r = yield d
        self.assertEqual(r, ['builder1', 'builder2'])
        self.protocol.get_message_result.assert_not_called()

    @defer.inlineCallbacks
    def test_remote_set_builder_list_key_is_missing(self) -> InlineCallbacksType[None]:
        d = self.set_up_set_builder_list([('builder1', 'test_dir1'), ('builder2', 'test_dir2')])
        self.check_message_set_worker_settings()

        yield self.check_message_send_response(
            'listdir', {'path': 'testdir'}, [('no_key', []), ('rc', 0)]
        )

        with self.assertRaisesRegex(Exception, "Key 'files' is missing."):
            yield d

        self.protocol.get_message_result.assert_not_called()

    @defer.inlineCallbacks
    def test_remote_set_builder_list_key_rc_not_zero(self) -> InlineCallbacksType[None]:
        d = self.set_up_set_builder_list([('builder1', 'test_dir1'), ('builder2', 'test_dir2')])
        self.check_message_set_worker_settings()

        yield self.check_message_send_response('listdir', {'path': 'testdir'}, [('rc', 123)])

        with self.assertRaisesRegex(Exception, "Error number: 123"):
            yield d

        self.protocol.get_message_result.assert_not_called()

    @parameterized.expand([
        ('want_stdout', 0, False),
        ('want_stdout', 1, True),
        ('want_stderr', 0, False),
        ('want_stderr', 1, True),
        (None, None, None),
    ])
    @defer.inlineCallbacks
    def test_remote_start_command_args_update(
        self,
        arg_name: str | None,
        arg_value: int | None,
        expected_value: bool | None,
    ) -> InlineCallbacksType[None]:
        self.protocol.get_message_result.return_value = defer.succeed(None)

        rc_instance = base.RemoteCommandImpl()
        result_command_id_to_command_map = {"1": rc_instance}
        self.protocol.command_id_to_command_map = {}
        args: dict[str, Any] = {'args': 'args'}

        if arg_name is not None:
            args[arg_name] = arg_value

        yield self.conn.remoteStartCommand(rc_instance, 'builder', "1", 'command', args)
        expected_args = args.copy()
        if arg_name is not None:
            expected_args[arg_name] = expected_value

        self.assertEqual(result_command_id_to_command_map, self.protocol.command_id_to_command_map)
        self.protocol.get_message_result.assert_called_with({
            'op': 'start_command',
            'builder_name': 'builder',
            'command_id': "1",
            'command_name': 'command',
            'args': expected_args,
        })

    @defer.inlineCallbacks
    def test_remote_shutdown(self) -> InlineCallbacksType[None]:
        self.protocol.get_message_result.return_value = defer.succeed(None)
        yield self.conn.remoteShutdown()

        self.protocol.get_message_result.assert_called_once_with({'op': 'shutdown'})

    @defer.inlineCallbacks
    def test_remote_interrupt_command(self) -> InlineCallbacksType[None]:
        self.protocol.get_message_result.return_value = defer.succeed(None)
        yield self.conn.remoteInterruptCommand('builder', 1, 'test')

        self.protocol.get_message_result.assert_called_once_with({
            'op': 'interrupt_command',
            'builder_name': 'builder',
            'command_id': 1,
            'why': 'test',
        })

    def test_perspective_keepalive(self) -> None:
        self.conn.perspective_keepalive()
        self.conn.worker.messageReceivedFromWorker.assert_called_once_with()  # type: ignore[attr-defined]

    def test_perspective_shutdown(self) -> None:
        self.conn.perspective_shutdown()

        self.conn.worker.shutdownRequested.assert_called_once_with()  # type: ignore[attr-defined]
        self.conn.worker.messageReceivedFromWorker.assert_called_once_with()  # type: ignore[attr-defined]
